import torch
from torch.utils.data import Dataset
import pandas as pd

# ================== 新增 ReturnsDataset 类 ==================
class ReturnsDataset(Dataset):
    """用于时间序列训练数据的 PyTorch Dataset"""
    def __init__(self, data: pd.DataFrame, window_size: int = 30):
        """
        参数:
            data (pd.DataFrame): 包含 `returns` 列的数据
            window_size (int): 输入序列的时间窗口长度
        """
        self.data = data["returns"].values  # 转换为 NumPy 数组
        self.window_size = window_size

    def __len__(self) -> int:
        return len(self.data) - self.window_size

    def __getitem__(self, idx: int) -> tuple:
        """
        返回:
            x (torch.Tensor): 输入序列 [window_size]
            y (torch.Tensor): 预测目标 (下一时刻的收益率)
        """
        x = self.data[idx:idx + self.window_size]
        y = self.data[idx + self.window_size]
        return (
            torch.tensor(x, dtype=torch.float32),
            torch.tensor(y, dtype=torch.float32)
        )
